package io.a.function;

import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.expressions.Aggregator;

public class AvgAgeAggregator extends Aggregator <Long, AvgAgeBuffer, Long> {

    @Override
    public AvgAgeBuffer zero() {
        return new AvgAgeBuffer(0L, 0L);
    }

    @Override
    public AvgAgeBuffer reduce(AvgAgeBuffer buffer, Long in) {
        buffer.setTotal(buffer.getTotal() + in);
        buffer.setCnt(buffer.getCnt() + 1);
        return buffer;
    }

    @Override
    public AvgAgeBuffer merge(AvgAgeBuffer buffer1, AvgAgeBuffer buffer2) {
        AvgAgeBuffer buffer = new AvgAgeBuffer(0L, 0L);
        buffer.setTotal(buffer1.getTotal() + buffer2.getTotal());
        buffer.setCnt(buffer1.getCnt() + buffer2.getCnt());
        return buffer;
    }

    @Override
    public Long finish(AvgAgeBuffer buffer) {
        return buffer.getTotal() / buffer.getCnt();
    }

    @Override
    public Encoder<AvgAgeBuffer> bufferEncoder() {
        return Encoders.bean(AvgAgeBuffer.class);
    }

    @Override
    public Encoder<Long> outputEncoder() {
        return Encoders.LONG();
    }
}
